import torch
import pickle
import matplotlib.pyplot as plt


def tf2(label):
    return torch.tensor(label)


# 1. 加载截取的数据集
with open('subset.pkl', 'rb') as f:
    ds = pickle.load(f)

# 2. 查看长度
print(len(ds))

# 3. 查看元素
print(ds[0])

# 4. 学生练习：创建数据加载器，按照batch_size = 32进行数据加载, 显示第一批次的6张图像
from torch.utils.data import DataLoader

train_loader = DataLoader(ds, batch_size=32, shuffle=True)
for img, label in train_loader:
    print(f"图像的大小 {img.shape} 对应的标签：{label.shape}")
    plt.figure(figsize=(10, 10))
    for i in range(6):
        print(img[i].shape, label[i])
        plt.subplot(2, 3, i + 1)
        # plt.imshow(img[i].squeeze().numpy(), cmap='gray')  # 单通道也可以直接压缩为二维显示
        plt.imshow(img[i].permute(1, 2, 0), cmap='gray')
        plt.title(label[i].item())
        plt.axis('off')
    break
plt.show()